import random
import numpy as np

from typing import Union, List
from case_generator.generator.generate_types import GENERATOR_REGISTRY
from case_generator.generator.base_generator import CaseGenerator
from configs.case_config import InputCaseConfig


@GENERATOR_REGISTRY.register("ascend_generate_aclnn_moe_token_permute")
class DefaultGenerator(CaseGenerator):
    def __init__(self, config):
        super().__init__(config)
        self.tokens_shape = []

    def after_input_config(
            self,
            index: int,
            input_case: Union[InputCaseConfig, List[InputCaseConfig]]
    ) -> Union[InputCaseConfig, List[InputCaseConfig]]:

        if input_case.name == "tokens":
            self.tokens_shape = input_case.shape

        if input_case.name == "indices":
            if self.tokens_shape == [10135, 5376]: ##新增现网问题用例，指定当tokens的shape大小为[10135, 5376]时，indices的shape大小为[10135]
                input_case.shape = [10135]
            else:
                input_case.shape[0] = self.tokens_shape[0]
                if len(input_case.shape) == 2:
                    input_case.shape[1] = random.randint(1, 512)

        return input_case